import numpy as np
import scipy.sparse as sp
from scipy.sparse.linalg import eigsh
import torch
from torch_geometric.data import Data
from tqdm import tqdm  # 引入 tqdm 进度条库

# ================== 安全加载配置 ==================
torch.serialization.add_safe_globals([Data])

# ================== 数据加载模块 ==================
def load_graph_data(graph_name):
    """通用数据加载器，支持多种存储格式"""
    # 加载原始数据
    loaded_data = torch.load(
        f"../data/{graph_name}_data.pt",
        weights_only=True,
        map_location=torch.device('cpu')
    )
    
    # 统一数据格式
    if isinstance(loaded_data, (tuple, list)):
        data = loaded_data[0]
    else:
        data = loaded_data
    
    # 转换为字典格式访问
    if isinstance(data, Data):
        data_dict = {
            'x': data.x,
            'edge_index': data.edge_index,
            'y': data.y
        }
    else:
        data_dict = data
    
    # 提取边索引
    edge_index = data_dict['edge_index'].numpy()
    num_nodes = data_dict['x'].shape[0]
    
    # 构建对称邻接矩阵
    row = edge_index[0]
    col = edge_index[1]
    adj = sp.coo_matrix(
        (np.ones_like(row), (row, col)),
        shape=(num_nodes, num_nodes),
        dtype=np.float32
    )
    adj = adj.maximum(adj.T)
    
    # 验证邻接矩阵
    print("\n邻接矩阵验证:")
    print(f"节点总数: {num_nodes}")
    print(f"有效边数量: {adj.count_nonzero()}")
    print(f"对称性检查: {(adj != adj.T).sum() == 0}")
    
    return adj, num_nodes

# ================== 特征分解核心 ==================
def laplacian_eigenmaps(adj, emb_dim=128):
    """保证输出维度稳定的特征分解实现"""
    # 计算归一化拉普拉斯矩阵
    degree = np.array(adj.sum(1)).flatten()
    D_inv_sqrt = sp.diags(np.power(degree, -0.5, where=degree != 0))
    L = sp.eye(adj.shape[0]) - D_inv_sqrt @ adj @ D_inv_sqrt
    
    # 动态调整特征分解参数
    max_possible_k = min(adj.shape[0] - 2, 300)  # ARPACK限制
    target_k = min(emb_dim + 20, max_possible_k)
    
    # 特征分解参数配置
    eig_params = {
        'k': target_k,
        'which': 'SM',
        'maxiter': 200000,
        'tol': 1e-12,
        'ncv': 500
    }
    
    # 执行分解
    eigenvalues, eigenvectors = eigsh(L, **eig_params)
    
    # 筛选有效特征向量
    valid_idx = np.where(eigenvalues > 1e-10)[0]
    valid_vectors = eigenvectors[:, valid_idx]
    
    # 维度保证机制
    if valid_vectors.shape[1] < emb_dim:
        print(f"警告：仅找到{valid_vectors.shape[1]}个有效向量，使用随机补充")
        missing = emb_dim - valid_vectors.shape[1]
        random_vectors = np.random.randn(adj.shape[0], missing)
        # 去除随机向量与已有特征向量的相关性
        random_vectors = random_vectors - valid_vectors @ (valid_vectors.T @ random_vectors)
        random_vectors /= np.linalg.norm(random_vectors, axis=0)
        valid_vectors = np.hstack([valid_vectors, random_vectors])
    
    return valid_vectors[:, :emb_dim]

# ================== 嵌入保存模块 ==================
def save_embeddings(embeddings, num_nodes, emb_dim, output_file):
    """带严格校验的保存函数"""
    # 维度验证
    assert embeddings.shape == (num_nodes, emb_dim), \
        f"维度不匹配！预期({num_nodes}, {emb_dim})，实际{embeddings.shape}"
    
    # 数值验证
    assert not np.isnan(embeddings).any(), "嵌入包含NaN值"
    assert np.isfinite(embeddings).all(), "嵌入包含无限值"
    
    # 写入文件
    with open(output_file, 'w') as f:
        f.write(f"{num_nodes}\t{emb_dim}\n")
        # 使用 tqdm 进度条显示保存进度
        for node_id in tqdm(range(num_nodes), desc="保存嵌入进度", unit="node"):
            vector = embeddings[node_id]
            line = [f"{node_id}"] + [f"{x:.6f}" for x in vector]
            f.write("\t".join(line) + "\n")

# ================== 主程序 ==================
if __name__ == "__main__":
    # 用户配置
    graph_name = input("请输入图名称（如Cora/Citeseer）: ").strip()
    EMB_DIM = 128
    OUTPUT_FILE = f"../data/{graph_name}_lap.emb"
    
    try:
        print("\n====== 数据加载阶段 ======")
        adj, num_nodes = load_graph_data(graph_name)
        
        print("\n====== 特征分解阶段 ======")
        print(f"目标维度: {EMB_DIM}")
        embeddings = laplacian_eigenmaps(adj, EMB_DIM)
        
        print("\n====== 嵌入验证阶段 ======")
        print(f"最终嵌入维度: {embeddings.shape[1]}")
        print(f"首节点样例: {embeddings[0][:5]}...")
        
        print("\n====== 文件保存阶段 ======")
        save_embeddings(embeddings, num_nodes, EMB_DIM, OUTPUT_FILE)
        print(f"成功生成: {OUTPUT_FILE}")
        
    except Exception as e:
        print(f"\n错误发生: {str(e)}")
        print("建议检查：")
        print("- 数据文件路径是否正确")
        print("- 图数据是否包含完整边信息")
        print("- 系统内存是否充足")
